import numpy as np
import netCDF4 as nc
import matplotlib.pyplot as plt
import cartopy.crs as ccrs
import cartopy.feature as cfeature
from matplotlib.colors import ListedColormap

# -------------------------- 1. 数据加载函数 --------------------------
def load_cloud_data(file_path=None):
    """加载 NetCDF 数据，文件不存在则生成模拟数据"""
    try:
        with nc.Dataset(file_path, 'r') as nf:
            clt_data = np.ma.getdata(nf.variables['clt'][:])
            cth_data = np.ma.getdata(nf.variables['cth'][:])
            cbh_data = np.ma.getdata(nf.variables['predicted_mean'][:])
            uncertainty_data = np.ma.getdata(nf.variables['predicted_uncertainty'][:])
            lat = nf.variables['lat'][:]
            lon = nf.variables['lon'][:]
        print(f"成功加载数据：{file_path}")
        return clt_data, cth_data, cbh_data, uncertainty_data, lat, lon
    except FileNotFoundError:
        print("警告：未找到指定文件，使用随机数据演示")
        # 生成模拟经纬度网格
        lon, lat = np.meshgrid(np.linspace(70, 140, 200), np.linspace(0, 60, 100))
        # 模拟云属性数据
        clt_data = np.random.randint(2, 8, size=lon.shape)
        cth_data = np.random.rand(*lon.shape) * 15
        cbh_data = cth_data - np.random.rand(*lon.shape) * 5
        uncertainty_data = np.random.rand(*lon.shape) * 4
        clt_data[clt_data < 2] = 2  # 保证云类型数据范围
        return clt_data, cth_data, cbh_data, uncertainty_data, lat, lon


# -------------------------- 2. 数据预处理函数 --------------------------
def preprocess_data(clt_data, cth_data, cbh_data, uncertainty_data, lat, lon):
    """处理无效值、过滤云底高度负值"""
    # 转换 masked 数组为带 NaN 的数组
    all_data = [clt_data, cth_data, cbh_data, uncertainty_data]
    clt_data, cth_data, cbh_data, uncertainty_data = [
        np.ma.masked_invalid(d).filled(np.nan) for d in all_data
    ]
    # 过滤云底高度小于 0 的值
    cbh_data[cbh_data < 0] = np.nan
    # 处理经纬度无效值（用均值填充）
    lon = np.nan_to_num(lon, nan=np.nanmean(lon))
    lat = np.nan_to_num(lat, nan=np.nanmean(lat))
    return clt_data, cth_data, cbh_data, uncertainty_data, lat, lon


# -------------------------- 3. 可视化绘图函数 --------------------------
def create_cloud_visualization(clt_data, cth_data, cbh_data, uncertainty_data, lat, lon, output_path):
    """创建 2x2 布局可视化，对换颜色条位置"""
    # 计算云高、不确定性的全局范围（用于共享颜色条）
    data_for_scaling = np.array([cth_data, cbh_data, uncertainty_data])
    global_min = np.nanmin(data_for_scaling)
    global_max = np.nanmax(data_for_scaling)
    print(f"CTH/CBH/Uncertainty 统一颜色范围: {global_min:.2f} - {global_max:.2f} km")

    # 云类型配置（名称 + 颜色映射）
    cloud_types = {
        2: "Water",
        3: "SuperCooled",
        4: "Mixed",
        5: "Ice",
        6: "Cirrus",
        7: "Overlap"
    }
    colors = {
        2: "#F9F8CA",
        3: "#96D2B0",
        4: "#35B9C5",
        5: "#2681B6",
        6: "#1E469B",
        7: "#080f40"
    }

    # 创建 2x2 子图 + 左右侧颜色条区域布局
    fig, axes = plt.subplots(
        nrows=2, ncols=2,
        figsize=(12, 12),  # 画布大小适配布局
        subplot_kw={'projection': ccrs.Orthographic(central_longitude=105, central_latitude=0)}
    )
    axes = axes.flatten()  # 转为一维数组方便遍历

    # 定义共享颜色映射（云高/不确定性用 turbo）
    # shared_cmap = 'turbo'
    # shared_cmap = 'rainbow'
    shared_cmap = 'jet'
    # shared_cmap = 'Blues'

    # -------------------- (1) 绘制云顶高度 (Cloud Top Height) --------------------
    ax_cth = axes[0]
    im_cth = ax_cth.pcolormesh(
        lon, lat, cth_data,
        transform=ccrs.PlateCarree(),
        cmap=shared_cmap, vmin=global_min, vmax=global_max
    )
    ax_cth.set_title('Cloud Top Height', fontsize=12)

    # -------------------- (2) 绘制云底高度 (Cloud Base Height) --------------------
    ax_cbh = axes[1]
    im_cbh = ax_cbh.pcolormesh(
        lon, lat, cbh_data,
        transform=ccrs.PlateCarree(),
        cmap=shared_cmap, vmin=global_min, vmax=global_max
    )
    ax_cbh.set_title('Cloud Base Height', fontsize=12)

    # -------------------- (3) 绘制云类型 (Cloud Type) --------------------
    ax_clt = axes[2]
    # 构建云类型离散颜色映射
    unique_clt = np.unique(clt_data[~np.isnan(clt_data)])
    valid_clt = [t for t in unique_clt if t in cloud_types]
    clt_cmap_list = [colors[t] for t in valid_clt]
    clt_cmap = ListedColormap(clt_cmap_list)
    im_clt = ax_clt.pcolormesh(
        lon, lat, clt_data,
        transform=ccrs.PlateCarree(),
        cmap=clt_cmap
    )
    ax_clt.set_title('Cloud Type', fontsize=12)

    # -------------------- (4) 绘制不确定性 (Uncertainty) --------------------
    ax_unc = axes[3]
    im_unc = ax_unc.pcolormesh(
        lon, lat, uncertainty_data,
        transform=ccrs.PlateCarree(),
        cmap=shared_cmap, vmin=global_min, vmax=global_max
    )
    ax_unc.set_title('Cloud Base Height Uncertainty', fontsize=12)

    # -------------------- 统一设置地图样式 --------------------
    for ax in axes:
        ax.set_global()  # 显示全球范围
        ax.add_feature(cfeature.COASTLINE.with_scale('50m'), linewidth=0.5, edgecolor='black')
        ax.add_feature(cfeature.BORDERS, linestyle=':', linewidth=0.4, edgecolor='black')
        ax.gridlines(draw_labels=False, color='black', alpha=0.3, linestyle='--')



    # -------------------- 调整子图间距，为底部颜色条预留空间 --------------------
    fig.subplots_adjust(bottom=0.12, wspace=0.1, hspace=0.15)

    # -------------------- 颜色条位置调整（底部并排） --------------------
    # 1. 云类型颜色条 → 左下横向
    cbar_clt_ax = fig.add_axes([0.15, 0.08, 0.35, 0.03])  # [左, 下, 宽, 高]
    cbar_clt = fig.colorbar(
        im_clt, cax=cbar_clt_ax, orientation='horizontal', ticks=valid_clt
    )
    cbar_clt.set_ticklabels([cloud_types[t] for t in valid_clt])
    cbar_clt.set_label('Cloud Type', fontsize=10)

    # 2. 云高/不确定性共享颜色条 → 右下横向，与云类型颜色条并排
    cbar_shared_ax = fig.add_axes([0.55, 0.08, 0.35, 0.03])  # [左, 下, 宽, 高]
    cbar_shared = fig.colorbar(
        im_cth, cax=cbar_shared_ax, orientation='horizontal'
    )
    cbar_shared.set_label('Height / Uncertainty (km)', fontsize=10)
    # # 调整子图间距，为底部颜色条预留空间
    # fig.subplots_adjust(bottom=0.18, wspace=0.3, hspace=0.2)
    #
    # # 1. 共享颜色条（CTH/CBH/Uncertainty）
    # shared_cbar_ax = fig.add_axes([0.2, 0.08, 0.35, 0.03])  # 位置：左、下、宽、高
    # shared_cbar = fig.colorbar(im1, cax=shared_cbar_ax, orientation='horizontal')
    # shared_cbar.set_label('Height / Uncertainty (km)', fontsize=12)
    #
    # # 2. 云类型颜色条，与共享条并排
    # clt_cbar_ax = fig.add_axes([0.6, 0.08, 0.35, 0.03])  # 并排摆放
    # clt_cbar = fig.colorbar(im3, cax=clt_cbar_ax, orientation='horizontal')
    # clt_cbar.set_label('Cloud Type Value', fontsize=12)

    # -------------------- 总标题与保存 --------------------
    fig.suptitle('Cloud Property Analysis (2020-06-27 05:00)', fontsize=16, y=0.95)
    plt.savefig(output_path, dpi=300, facecolor='white', bbox_inches='tight')
    print(f"图像已保存至: {output_path}")
    return fig


# -------------------------- 4. 主程序入口 --------------------------
if __name__ == "__main__":
    # 替换为实际文件路径
    file_path = '/mnt/datastore/liudddata/result/20200506_droupout/2020052105_predicted_2d_mc.nc'
    output_path = '/mnt/datastore/liudddata/cloudsat_data/cloudsat_cbh_csv/cloud_analysis_final.png'

    # 数据流程：加载 → 预处理 → 可视化
    clt, cth, cbh, unc, lat, lon = load_cloud_data(file_path)
    clt, cth, cbh, unc, lat, lon = preprocess_data(clt, cth, cbh, unc, lat, lon)
    create_cloud_visualization(clt, cth, cbh, unc, lat, lon, output_path)
    plt.show()